import psutil
import os
import shutil
import time
import threading
from pystray import Icon, MenuItem, Menu
from PIL import Image, ImageDraw
from watchdog.observers import Observer
from watchdog.events import FileSystemEventHandler
from concurrent.futures import ThreadPoolExecutor
import logging

# 设置日志记录
logging.basicConfig(filename="usb_format_log.txt", level=logging.INFO, format="%(asctime)s - %(message)s")

# 用于存储已经检测到的设备编号
detected_devices = set()
# 状态标识，控制脚本是否执行
script_running = True


# 删除USB设备内所有文件的函数
def delete_files_in_directory(directory):
    try:
        # 遍历目录并删除文件和子目录
        for root, dirs, files in os.walk(directory, topdown=False):
            for name in files:
                file_path = os.path.join(root, name)
                try:
                    os.remove(file_path)
                    print(f"已删除文件：{file_path}")
                    logging.info(f"已删除文件：{file_path}")
                except Exception as e:
                    print(f"删除文件失败：{file_path}, 错误：{e}")
                    logging.error(f"删除文件失败：{file_path}, 错误：{e}")
            for name in dirs:
                dir_path = os.path.join(root, name)
                try:
                    os.rmdir(dir_path)
                    print(f"已删除目录：{dir_path}")
                    logging.info(f"已删除目录：{dir_path}")
                except Exception as e:
                    print(f"删除目录失败：{dir_path}, 错误：{e}")
                    logging.error(f"删除目录失败：{dir_path}, 错误：{e}")
    except Exception as e:
        print(f"删除目录 {directory} 中的文件时发生错误：{e}")
        logging.error(f"删除目录 {directory} 中的文件时发生错误：{e}")


# 文件系统事件处理类，用于监测新文件操作
class FileSystemHandler(FileSystemEventHandler):
    def on_created(self, event):
        if event.is_directory:
            return
        print(f"检测到新文件：{event.src_path}")
        # 文件创建后立即删除
        try:
            os.remove(event.src_path)
            print(f"已删除新文件：{event.src_path}")
            logging.info(f"已删除新文件：{event.src_path}")
        except Exception as e:
            print(f"删除新文件失败：{event.src_path}, 错误：{e}")
            logging.error(f"删除新文件失败：{event.src_path}, 错误：{e}")


# 检查并检测插入的 USB 设备
def check_for_usb():
    while True:
        if script_running:
            # 获取当前系统中的所有磁盘设备
            partitions = psutil.disk_partitions()
            for partition in partitions:
                # 过滤出 USB 设备，通常以 "removable" 标识
                if 'removable' in partition.opts:
                    device = partition.device
                    if device not in detected_devices:
                        detected_devices.add(device)
                        current_time = time.strftime("%Y-%m-%d %H:%M:%S", time.localtime())
                        print(f"检测到新的USB设备：{device} at {current_time}")
                        logging.info(f"检测到新的USB设备：{device} at {current_time}")
                        # 启动并发操作删除 USB 设备中的所有文件
                        executor.submit(delete_files_in_directory, partition.mountpoint)

                        # 设置文件监视器
                        event_handler = FileSystemHandler()
                        observer = Observer()
                        observer.schedule(event_handler, partition.mountpoint, recursive=True)
                        observer.start()
                        print(f"开始监控 {partition.mountpoint} 的文件操作")
                        logging.info(f"开始监控 {partition.mountpoint} 的文件操作")

                    else:
                        # 检测到 USB 拔出
                        current_time = time.strftime("%Y-%m-%d %H:%M:%S", time.localtime())
                        print(f"USB设备拔出：{device} at {current_time}")
                        logging.info(f"USB设备拔出：{device} at {current_time}")
                        detected_devices.remove(device)

        time.sleep(3)  # 每3秒检测一次


# 创建系统托盘图标并提供菜单
def create_tray_icon():
    # 创建图标
    image = Image.new('RGB', (64, 64), color=(255, 255, 255))
    draw = ImageDraw.Draw(image)
    draw.rectangle([0, 0, 64, 64], fill=(0, 128, 255))  # 设置图标的颜色

    # 定义菜单项
    def on_quit(icon, item):
        global script_running
        script_running = False  # 停止脚本的执行
        icon.stop()
        print("退出程序")

    def on_toggle(icon, item):
        global script_running
        script_running = not script_running  # 切换脚本的执行状态
        status = "已启动" if script_running else "已暂停"
        print(f"脚本执行状态：{status}")
        item.text = f"暂停执行" if script_running else f"恢复执行"  # 更改菜单文本

    # 创建菜单
    menu = Menu(MenuItem("切换脚本执行", on_toggle), MenuItem("退出", on_quit))

    # 创建并运行图标
    icon = Icon("USB Auto Formatter", image, menu=menu)
    icon.run()


# 主程序
if __name__ == '__main__':
    # 创建并启动系统托盘图标线程
    tray_thread = threading.Thread(target=create_tray_icon, daemon=True)
    tray_thread.start()

    # 创建并启动并发执行器（ThreadPoolExecutor）
    executor = ThreadPoolExecutor(max_workers=5)

    # 启动检测USB设备的线程
    check_for_usb()
